from selenium import webdriver
from time import sleep, time
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.keys import Keys
import requests
import request
import uuid
import fileUtils
import cvUtils
import cv2
import queryChromeHistory

import os
from selenium.webdriver.chrome.options import Options


def startFetch(url):
    """开始抓取脚本"""
    # os.system("taskkill /f /im chrome.exe")

    # 设置下载路径
    download_path = r"H:\download\tmp"
    chrome_history_path = r"H:\download\chrome\profile"
    # 设置 ChromeOptions 来配置下载路径和自动下载
    chrome_options = Options()
    chrome_options.add_experimental_option("prefs", {
        "download.default_directory": download_path,
        "download.prompt_for_download": False,
        "download.directory_upgrade": True,
        "safebrowsing.enabled": True,
        "profile.default_content_settings.popups": 0
    })
    # 指定用户数据目录
    chrome_options.add_argument(r"--user-data-dir=H:\download\userdata")
    # 下载文件数据库的保存地址
    chrome_options.add_argument(r"--user-data-dir=" + chrome_history_path)
    chrome_options.add_argument("--start-minimized")  # 设置窗口默认最小化

    # 1.创建Chrome浏览器对象，这会在电脑上在打开一个浏览器窗口
    service = Service('D:\chromedriver-win64\chromedriver.exe')
    driver = webdriver.Chrome(service=service, options=chrome_options)

    # 4.设置浏览器的大小
    driver.set_window_size(1200, 700)
    # 2.通过浏览器向服务器发送URL请求
    driver.get(url)
    sleep(1)

    # 判断是否需要登录才能访问
    needLogin = False
    try:
        loginElement = driver.find_element(By.XPATH, "//*[contains(text(), 'Log in to view')]")
        print("需要登录才能访问，等待操作登录")
        needLogin = True
        sleep(1)
        return "", "", "需要登录才能访问", "", "", "", "需要登录才能访问"
    except:
        print("无需登录就可以访问")
        sleep(1)
        needLogin = False

    # 判断是否登录
    while True:
        needLogin2 = False
        try:
            headerElement = driver.find_element(By.CSS_SELECTOR,
                                                "header.mantine-Header-root")
            loginElement = headerElement.find_element(By.XPATH, "//*[contains(text(), 'Sign In')]")
            print("找到sign In")
            needLogin2 = True
            sleep(1)
        except:
            sleep(1)

        if (needLogin2 is not True):
            try:
                stackElement = driver.find_element(By.CSS_SELECTOR,
                                                   "header.mantine-Stack-root")
                loginElement = stackElement.find_element(By.XPATH,
                                                         "//*[contains(text(), 'GitHub')]")
                print("准备登录github")
                needLogin2 = True
                sleep(1)
            except:
                sleep(1)

        if (needLogin2 is not True):
            try:
                stackElement = driver.find_element(By.CSS_SELECTOR,
                                                   "div.auth-form-body")
                loginElement = stackElement.find_element(By.XPATH,
                                                         "//*[contains(text(), 'GitHub')]")
                print("在github登录页")
                needLogin2 = True
                sleep(1)
            except:
                sleep(1)

        if needLogin2 is not True:
            print("已登录")
            break
        else:
            print("未登录,等待登录后再试")

    # 使用CSS选择器找到h1元素
    # 查找标题
    h1_element = driver.find_element(By.CSS_SELECTOR, "h1.mantine-Text-root")
    # 获取h1元素的文本值
    nameOrigin = h1_element.text

    # 遍历表格根据名称查询
    tableElement = driver.find_element(By.CSS_SELECTOR,
                                       "table.mantine-Table-root tbody")

    # 在div元素下查找所有的a标签
    type = ""
    version = ""
    triggerWord = ""
    imgUrl = ""
    trElements = tableElement.find_elements(By.TAG_NAME, 'tr')
    for element in trElements:
        tdElement = element.find_element(By.CSS_SELECTOR, 'td')

        # 查找lora和ckpt类型
        if "Type" in tdElement.text:
            tdElement2 = element.find_element(By.CSS_SELECTOR, 'td:nth-child(2)')
            type = tdElement2.text
        # 查找基础模型
        if "Base Model" in tdElement.text:
            tdElement2 = element.find_element(By.CSS_SELECTOR, 'td:nth-child(2)')
            version = tdElement2.text
        # 查找关键词
        if "Trigger" in tdElement.text:
            tdElement2 = element.find_element(By.CSS_SELECTOR,
                                              'td:nth-child(2) div.mantine-Group-root')

            wordsElement = tdElement2.find_elements(By.XPATH, './div')

            for e in wordsElement:
                triggerWord = triggerWord + e.text + ","

    try:
        # 查找首图
        imgElement = driver.find_element(By.CSS_SELECTOR,
                                         "div.mantine-Indicator-root div div.mantine-AspectRatio-root img")
        imgUrl = imgElement.get_attribute('src')
    except:
        print("抓取图片失败")

    # 查找下载地址
    downloadUrl = ""
    downloadElement = None
    try:
        print("开始第一种抓取下载地址的方式(需要点开下载按钮进行搜索的)")
        # 首先查找点击显示下载按钮的方式
        downloadIcon = driver.find_element(By.CSS_SELECTOR,
                                           "button.mantine-UnstyledButton-root div.mantine-Button-inner svg.tabler-icon-download")

        # 找到该元素的父元素
        downloadIcon = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
        downloadIcon = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
        downloadElement = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级

        downloadElement.click()
        sleep(0.5)
        downloadIcon = driver.find_element(By.CSS_SELECTOR,
                                           "div.mantine-Menu-dropdown div a.mantine-Menu-item div.mantine-Menu-itemIcon div.mantine-Group-root")
        downloadIcon = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
        downloadIcon = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
        downloadIcon = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
        parentElement = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级

        # 在div元素下查找所有的a标签
        a_elements = parentElement.find_elements(By.TAG_NAME, 'a')

        for element in a_elements:
            print(element.text)
            if "SafeTensor" in element.text \
                    or "LoRA" in element.text \
                    or "Embedding" in element.text:
                print("找到元素", element.text)
                downloadUrl = element.get_attribute('href')
                downloadElement = element

        # if downloadUrl is None or downloadUrl == "":
        #     raise Exception("None url")
    except:
        print("开始第二种抓取下载地址的方式,直接点击下载")
        downloadIcon = driver.find_element(By.CSS_SELECTOR,
                                           "a.mantine-UnstyledButton-root div.mantine-Button-inner svg.tabler-icon-download")
        downloadIcon = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
        downloadIcon = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
        downloadElement = downloadIcon.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级

        downloadUrl = downloadElement.get_attribute('href')

    print("下载地址为", downloadUrl)
    # downloadUrl = request.getDownloadUrl(downloadUrl)

    if downloadElement is not None:
        downloadElement.click()
        sleep(2)

    driver.quit()  # 退出整个WebDriver会话，关闭所有窗口
    sleep(1)
    if downloadElement is not None:
        localFilePath, url = queryChromeHistory.getLatestDownload(
            chrome_history_path=chrome_history_path + r"\Default\History")
        sleep(4)
        localFilePath, url = queryChromeHistory.getLatestDownload(
            chrome_history_path=chrome_history_path + r"\Default\History")
        downloadUrl = url

    # 下载图片
    # 生成一个UUID
    new_uuid = uuid.uuid4()
    # 将UUID转换为字符串
    uuid_str = str(new_uuid)
    suffix = fileUtils.getFileNameSuffix(imgUrl)
    localImgPath = download_path + "\\" + uuid_str + "." + suffix
    request.download_file(imgUrl, localImgPath)

    # 打印文本值
    print("头图地址", imgUrl)
    print("头图本地地址", localImgPath)
    print("标题为:", nameOrigin)
    print("类型为:", type)
    print("基础模型:", version)
    print("关键词:", triggerWord)
    print("最终下载地址为：", downloadUrl)

    return imgUrl, localImgPath, nameOrigin, type, version, triggerWord, downloadUrl


if __name__ == '__main__':
    # startFetch("https://civitai.com/models/84728/photon")
    # startFetch("https://civitai.com/models/7371/rev-animated")  # 需要点击两次
    # startFetch("https://civitai.com/models/65028/dwemer-tech-world-morph?modelVersionId=69655")
    # startFetch("https://civitai.com/models/6755/cetus-mix")
    # startFetch("https://civitai.com/models/94087/anxiang-or")  # 需要登录
    # startFetch("https://civitai.com/models/142656/ads-tech-world-morph")  # 需要登录
    # startFetch("https://civitai.com/models/200043?modelVersionId=262498")
    # startFetch(
    #     "https://civitai.com/models/135710/heterochromia-helper12-lora-pack?modelVersionId=149657")  # 无法下载的
    # startFetch(
    #     "https://civitai.com/models/124347?modelVersionId=152309")  # 下载失败的
    startFetch(
        "https://civitai.com/models/143636?modelVersionId=216791")
